import xml.dom.minidom as xmldom
import os

"""
 yolo5 将xml格式的文件转化为txt文件格式
"""
# 类型
classes = ['dljss', 'kcs', 'lmpss']


def data2txt(d1, d2, d3, d4, clsName, txtFile_name):
    fd = open(txtFile_name, 'a')
    index = classes.index(clsName)
    fd.write(str(index) + " " + str(d1) + " " + str(d2) + " " + str(d3) + " " + str(d4) + "\n")
    fd.close()


# 读取xml文件数据
def xml2data(xmlfile_name, txtFile_name):
    ele = xmldom.parse(os.path.abspath(xmlfile_name)).documentElement

    size = ele.getElementsByTagName("size")
    if len(size) != 0:
        image_width = int(size[0].getElementsByTagName("width")[0].firstChild.data)
        image_height = int(size[0].getElementsByTagName("height")[0].firstChild.data)

    objList = ele.getElementsByTagName("object")
    for obj in objList:
        index = objList.index(obj)
        clsName = objList[index].getElementsByTagName("name")[0].firstChild.data
        bndbox = objList[index].getElementsByTagName("bndbox")
        xmin = int(bndbox[0].getElementsByTagName("xmin")[0].firstChild.data)
        ymin = int(bndbox[0].getElementsByTagName("ymin")[0].firstChild.data)
        xmax = int(bndbox[0].getElementsByTagName("xmax")[0].firstChild.data)
        ymax = int(bndbox[0].getElementsByTagName("ymax")[0].firstChild.data)

        xcenter = (xmax + xmin) / 2
        ycenter = (ymax + ymin) / 2
        width = abs(xmax - xmin)
        height = abs(ymax - ymin)

        data2txt(xcenter / image_width, ycenter / image_height, width / image_width, height / image_height, clsName,
                 txtFile_name)


def readFile(filePath):
    xmlDir = filePath + "\\" + 'xml'
    if not os.path.exists(xmlDir):
        return None
    txtDir = filePath + "\\" + 'txt'
    if not os.path.exists(txtDir):
        os.mkdir(txtDir)
    for fileName in os.listdir(xmlDir):
        xml2data(xmlDir + '\\' + fileName, txtDir + '\\' + fileName[:-4] + ".txt")


if __name__ == '__main__':
    readFile("G:\\")
